#include "ip.h"

void handle_arp(struct rte_arp_hdr *arp_head, pSocketData d);
void handle_ipv4_icmp(struct rte_mbuf *_mbuf, pSocketData d);

struct rte_mbuf *rte_mbuf_deep_copy(struct rte_mbuf *src, struct rte_mempool *mp) {
	if (src == NULL || mp == NULL) {
		return NULL;
	}

	struct rte_mbuf *dst = rte_pktmbuf_alloc(mp);
	if (dst == NULL) {
		return NULL;
	}

	dst->pkt_len = src->pkt_len;
	dst->data_len = src->data_len;
	dst->nb_segs = src->nb_segs;
	dst->port = src->port;

	memcpy(rte_pktmbuf_mtod(dst, void *), rte_pktmbuf_mtod(src, void *), src->data_len);

	return dst;
}

static uint16_t checksum(uint16_t *addr, int count) {
	uint32_t sum = 0;
	uint16_t res = 0;

	while (count > 1) {
		sum += *(uint16_t *)addr++;
		count -= 2;
	}

	if (count > 0) {
		sum += *(uint8_t  *)addr;
	}

	while (sum >> 16) {
		sum = (sum & 0xffff) + (sum >> 16);
	}
	//res = sum & 0xffff;
	res = ~sum;

	return res;
}


uint16_t handle_ip(struct rte_mbuf *mbuf, pSocketData d)
{
	struct rte_ether_hdr *eth_head;
	struct rte_ipv4_hdr *ip_head;
	struct rte_udp_hdr *udp_head;
	struct rte_arp_hdr *arp_head;
	uint16_t data_len = 0;

	eth_head = rte_pktmbuf_mtod(mbuf, struct rte_ether_hdr*);
	switch (eth_head->ether_type) {
		case rte_cpu_to_be_16(RTE_ETHER_TYPE_ARP):
			arp_head = (struct rte_arp_hdr *)(eth_head + 1);
			handle_arp(arp_head, d);
			break;
		case rte_cpu_to_be_16(RTE_ETHER_TYPE_IPV4):
			ip_head = rte_pktmbuf_mtod_offset(mbuf, struct rte_ipv4_hdr*, sizeof(struct rte_ether_hdr));
			switch (ip_head->next_proto_id)
			{
				case IPPROTO_ICMP:
					handle_ipv4_icmp(mbuf, d);
					break;
				case IPPROTO_UDP:
					data_len = mbuf->data_len;
					break;
			}
			break;
		default:
			break;
	}
	return data_len;
}

void handle_arp(struct rte_arp_hdr *arp_head, pSocketData d) 
{
	if (arp_head->arp_data.arp_tip != d->local_ip) {
		return;
	}
	
	const uint16_t total_len = sizeof(struct rte_ether_hdr) + sizeof(struct rte_arp_hdr);
	struct rte_mbuf *mbuf = rte_pktmbuf_alloc(d->mpool); 
	if (mbuf == NULL) {
		fprintf(stderr, "Failed to allocate mbuf for ARP\n");
	}
	mbuf->pkt_len = total_len;
	mbuf->data_len = total_len;

	uint8_t *pktdata = rte_pktmbuf_mtod(mbuf, uint8_t*); 
	struct rte_ether_hdr *eth = (struct rte_ether_hdr *)pktdata;
#if RTE_VERSION == RTE_VERSION_NUM(19, 11, 0, 0)
	rte_memcpy(eth->s_addr.addr_bytes, d->local_mac, RTE_ETHER_ADDR_LEN);
	rte_memcpy(eth->d_addr.addr_bytes, arp_head->arp_data.arp_sha.addr_bytes, RTE_ETHER_ADDR_LEN);
#else
	rte_memcpy(eth->src_addr.addr_bytes, d->local_mac, RTE_ETHER_ADDR_LEN);
	rte_memcpy(eth->dst_addr.addr_bytes, arp_head->arp_data.arp_sha.addr_bytes, RTE_ETHER_ADDR_LEN); 
#endif
	eth->ether_type = htons(RTE_ETHER_TYPE_ARP);  // 设置以太类型为ARP

	struct rte_arp_hdr *arp = (struct rte_arp_hdr *)(eth + 1);
	arp->arp_hardware = htons(1);  // 硬件类型，1表示以太网
	arp->arp_protocol = htons(RTE_ETHER_TYPE_IPV4);  // 协议类型，IPv4
	arp->arp_hlen = RTE_ETHER_ADDR_LEN;  // 硬件地址长度
	arp->arp_plen = sizeof(uint32_t);  // 协议地址长度
	arp->arp_opcode = htons(RTE_ARP_OP_REPLY);  // 设置为ARP应答
	rte_memcpy(arp->arp_data.arp_sha.addr_bytes, d->local_mac, RTE_ETHER_ADDR_LEN);
	rte_memcpy(arp->arp_data.arp_tha.addr_bytes, arp_head->arp_data.arp_sha.addr_bytes, RTE_ETHER_ADDR_LEN);
	arp->arp_data.arp_sip = arp_head->arp_data.arp_tip;
	arp->arp_data.arp_tip = arp_head->arp_data.arp_sip; 

	rte_eth_tx_burst(d->portid, 0, &mbuf, 1); 
	rte_pktmbuf_free(mbuf); 
}


void handle_ipv4_icmp(struct rte_mbuf *_mbuf, pSocketData d)
{
	struct rte_ether_hdr *eth_head = rte_pktmbuf_mtod(_mbuf, struct rte_ether_hdr*);
	struct rte_ipv4_hdr *ip_head = (struct rte_ipv4_hdr *)(eth_head + 1);
	struct rte_icmp_hdr *icmp_head = (struct rte_icmp_hdr *)(ip_head + 1);
	
	struct rte_ether_hdr *eth;
	struct rte_ipv4_hdr *ip;
	struct rte_icmp_hdr *icmp;
	struct rte_mbuf *mbuf;
	static const uint32_t total_len = sizeof(struct rte_ether_hdr) + sizeof(struct rte_ipv4_hdr) + sizeof(struct rte_icmp_hdr);

	if (icmp_head->icmp_type == RTE_IP_ICMP_ECHO_REQUEST) {
		//mbuf = rte_pktmbuf_alloc(d->mpool);
		//mbuf->data_len = total_len;
		//mbuf->pkt_len = total_len;
		mbuf = rte_mbuf_deep_copy(_mbuf, d->mpool);
		eth = rte_pktmbuf_mtod(mbuf, struct rte_ether_hdr*);
		ip = rte_pktmbuf_mtod_offset(mbuf, struct rte_ipv4_hdr*, sizeof(struct rte_ether_hdr));
		icmp = (struct rte_icmp_hdr *)(ip + 1);

#if RTE_VERSION == RTE_VERSION_NUM(19, 11, 0, 0)
		rte_memcpy(eth->d_addr.addr_bytes, eth_head->s_addr.addr_bytes, RTE_ETHER_ADDR_LEN);
		rte_memcpy(eth->s_addr.addr_bytes, d->local_mac, RTE_ETHER_ADDR_LEN);
#else
		rte_memcpy(eth->dst_addr.addr_bytes, eth_head->src_addr.addr_bytes, RTE_ETHER_ADDR_LEN);
		rte_memcpy(eth->src_addr.addr_bytes, d->local_mac, RTE_ETHER_ADDR_LEN);
#endif
		eth->ether_type = htons(RTE_ETHER_TYPE_IPV4);
		
		ip->version_ihl = 0x45;
		ip->type_of_service = 0;
		//ip->total_length = htons(sizeof(struct rte_icmp_hdr) + sizeof(struct rte_ipv4_hdr));
		ip->packet_id = 0;
		ip->fragment_offset = 0;
		ip->time_to_live = 64;
		ip->next_proto_id = IPPROTO_ICMP;
		ip->dst_addr = ip_head->src_addr;
		ip->src_addr = d->local_ip;
		ip->hdr_checksum = 0;
		ip->hdr_checksum = rte_ipv4_cksum(ip);

		icmp->icmp_type = RTE_IP_ICMP_ECHO_REPLY;
		icmp->icmp_code = 0;
		icmp->icmp_ident = icmp_head->icmp_ident;
		icmp->icmp_seq_nb = icmp_head->icmp_seq_nb;
		icmp->icmp_cksum = 0;
		icmp->icmp_cksum = checksum((uint16_t *)icmp, ip->total_length-sizeof(struct rte_ipv4_hdr));

		int ret = rte_eth_tx_burst(d->portid, 0, &mbuf, 1);
		rte_pktmbuf_free(mbuf);

	}
}
